#include "complete.h"
#include <sstream>
#include <iostream>

namespace ads {
PrefixTree::PrefixTree() : value(-1) {}

PrefixTree::~PrefixTree() {
  for (tr1::unordered_map<char, PrefixTree*>::const_iterator cit = nodes.begin(); cit != nodes.end(); ++cit) {
    delete cit->second;
  }
}

void PrefixTree::insert(PrefixTree* root, const string& str, size_t pos) {
  if (pos < str.size()) {
    tr1::unordered_map<char, PrefixTree*>::const_iterator cit = root->nodes.find(str.at(pos));
    if (cit != root->nodes.end()) {
      insert(cit->second, str, pos + 1);
    } else {
      PrefixTree* pt = new PrefixTree();
      root->nodes[str.at(pos)] = pt;
      if (pos == (str.size() - 1)) {
        pt->value = 0;
      }
      insert(pt, str, pos + 1);
    }
  }
}

void PrefixTree::insert(const string& str) {
  insert(this, str, 0);
}


void PrefixTree::complete(PrefixTree* root, const string& pre, size_t pos, string newstr, multimap<int, string>* rest) {
  if (pos < pre.size()) {
    // searching
    tr1::unordered_map<char, PrefixTree*>::const_iterator cit = root->nodes.find(pre.at(pos));
    if (cit != root->nodes.end()) {
      complete(cit->second, pre, pos + 1, newstr, rest);
    }
  } else {
    // Fill all
    if (root->value >= 0) {
      rest->insert(pair<int,string>(root->value, newstr));
      root->value++;
    }
    for (tr1::unordered_map<char, PrefixTree*>::const_iterator cit = root->nodes.begin(); cit != root->nodes.end(); ++cit) {
      stringstream connectionstream;
      connectionstream << cit->first;
      complete(cit->second, pre, pos, newstr + connectionstream.str(), rest);
    }
  }
}

void PrefixTree::complete(const string& pre, vector<string>* rest) {
  multimap<int, string> rest_order;
  complete(this, pre, 0, "", &rest_order);
  for (multimap<int,string>::reverse_iterator it = rest_order.rbegin(); it != rest_order.rend(); ++it) {
    rest->push_back(it->second);
  }
}

bool PrefixTree::mark_remove(PrefixTree* root, const string& pre, size_t pos) {
  if (pos <= pre.size()) {
    if (pos == pre.size()) {
      root->value = -1;
      return true;
    } else {
      tr1::unordered_map<char,PrefixTree*>::const_iterator cit = root->nodes.find(pre.at(pos));
      if (cit != root->nodes.end()) {
        return mark_remove(cit->second, pre, pos + 1);
      }
    }
  }
  return false;
}

// Scan all the way till down, then keep on deleting and going up
bool PrefixTree::mark_remove(const string& pre) {
  return mark_remove(this, pre, 0);
}

}
